package com.ai.study.dashscope.controller.rag.milvus.utils;

import com.ai.study.dashscope.utils.PropertyConfigureUtil;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.DataType;
import io.milvus.param.ConnectParam;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.collection.*;
import io.milvus.param.index.CreateIndexParam;
import io.swagger.v3.oas.annotations.tags.Tag;
import lombok.extern.slf4j.Slf4j;

/**
 * @Author chenwei
 * @Date 2025/5/29 9:53
 */
@Slf4j
@Tag(name = "Milvus工具类")
public class MivusUtil {

    static String COLLECTION_NAME = "vector_store_default";

    // 读取配置属性
    static {
        COLLECTION_NAME = PropertyConfigureUtil.getProperty("spring.ai.vectorstore.milvus.collection-name");
    }

    /**
     * 获取milvus连接客户端
     * @return
     */
    public static MilvusServiceClient getMilvusClient() {
        MilvusServiceClient milvusClient = new MilvusServiceClient(
                ConnectParam.newBuilder()
                        .withHost("192.168.0.64")
                        .withPort(19530)
                        .build()
        );

        return milvusClient;
    }

    /**
     * 删除集合
     * @param milvusClient
     */
    public static void deleteCollection(MilvusServiceClient milvusClient) {

        if (!isCollectionExists(milvusClient)){
            log.info("集合不存在");
            return;
        }

        // 从内存中释放Collection
        milvusClient.releaseCollection(
                ReleaseCollectionParam.newBuilder()
                        .withCollectionName(COLLECTION_NAME)
                        .build()
        );
        // 删除Collection
        milvusClient.dropCollection(
                DropCollectionParam.newBuilder()
                        .withCollectionName(COLLECTION_NAME)
                        .build()
        );
    }

    /**
     * 判断集合是否存在
     * @param milvusClient
     * @return
     */
    public static Boolean isCollectionExists(MilvusServiceClient milvusClient){
        R<Boolean> respHasCollection = milvusClient.hasCollection(
                HasCollectionParam.newBuilder()
                        .withCollectionName(COLLECTION_NAME)
                        .build()
        );
        if (respHasCollection.getData().equals(Boolean.TRUE)) {
            return true;
        }
        return false;
    }

    /**
     * 创建集合
     * @param milvusClient
     */
    public static void createCollection(MilvusServiceClient milvusClient) {
        // 主键字段
        FieldType fieldType1 = FieldType.newBuilder()
                .withName("doc_id")
                .withDataType(DataType.VarChar)
                .withMaxLength(128)
                .withPrimaryKey(true)
                .withAutoID(false)
                .build();
        // 标量字段
        FieldType fieldType2 = FieldType.newBuilder()
                .withName("content")
                .withDataType(DataType.VarChar)
                .withMaxLength(256)
                .build();
        FieldType fieldType3 = FieldType.newBuilder()
                .withName("metadata")
                .withDataType(DataType.JSON)
                .build();
        // 二维浮点向量字段
        FieldType fieldType4 = FieldType.newBuilder()
                .withName("embedding")
                .withDataType(DataType.FloatVector)
                .withDimension(1536)
                .build();
        CreateCollectionParam createCollectionReq = CreateCollectionParam.newBuilder()
                .withCollectionName(COLLECTION_NAME)
                .withShardsNum(2)
                .addFieldType(fieldType1)
                .addFieldType(fieldType2)
                .addFieldType(fieldType3)
                .addFieldType(fieldType4)
                .build();
        milvusClient.createCollection(createCollectionReq);

        // 创建索引
        milvusClient.createIndex(
                CreateIndexParam.newBuilder()
                        .withCollectionName(COLLECTION_NAME)
                        .withFieldName("doc_id")
                        .withIndexType(IndexType.AUTOINDEX)
                        .build()
        );
        milvusClient.createIndex(
                CreateIndexParam.newBuilder()
                        .withCollectionName(COLLECTION_NAME)
                        .withFieldName("content")
                        .withIndexType(IndexType.AUTOINDEX)
                        .build()
        );
        milvusClient.createIndex(
                CreateIndexParam.newBuilder()
                        .withCollectionName(COLLECTION_NAME)
                        .withFieldName("embedding")
                        .withIndexType(IndexType.IVF_FLAT)
                        .withMetricType(MetricType.COSINE)
                        .withExtraParam("{\"nlist\":128}")
                        .withSyncMode(Boolean.FALSE)
                        .build()
        );

        // 加载Collection到内存
        milvusClient.loadCollection(
                LoadCollectionParam.newBuilder()
                        .withCollectionName(COLLECTION_NAME)
                        .build()
        );
    }
}
